using DataFrames
using QuadGK
using Distributions
using CSV
using LinearAlgebra
using JLD
using Interpolations
using DelimitedFiles

clearconsole()

#-------------------------------------------------------------
# IRF Configuration
#-------------------------------------------------------------

H = 60 # maximum horizon
sh_size = 3   # shock size, in multiples of standard deviations
n_drawsread = 200 #100 # transform draws 1 to n_drawsread

#-------------------------------------------------------------
# include Functions
#-------------------------------------------------------------
#cd("$(pwd())/Dropbox/Heterogeneity/Software/KS_Simulation/")
readDir = "$(pwd())/Functions/"
include(readDir *"vech.jl");
include(readDir *"logSpline_Procedures.jl");
include(readDir *"VAR_Procedures.jl");
include(readDir *"Loaddata.jl");
include(readDir *"EmpPercentiles_Procedures.jl")
include(readDir *"GiniFracBelowCutoff_Procedures.jl")

#-------------------------------------------------------------
# choose specification files
#-------------------------------------------------------------
nDataSel  = "2"
nfVARSpec = "3"
nModSpec  = "4"
nMCMCSpec = "1"
modName   = "VAR"  # VAR or SS

#cutoff    = 1;
nShockspec= "AggSh1"
# Aggregate shocks 1: TFP shock

specDir   = "$(pwd())/SpecFiles/"
include(specDir * "/fVARspec" * nfVARSpec * ".jl")
include(specDir * "/" * modName * "spec" * nModSpec * ".jl")
include(specDir * "/" * modName * "MCMCspec" * nMCMCSpec * ".jl")

# Percentiles of interest
ngrid     = 20
grid_temp = range(0.2, stop=2, length=ngrid);
grid_temp = [0.0; grid_temp]
vec_percs = [0.1; 0.2; 0.5; 0.8; 0.9]

cutoff = 2.0

#-------------------------------------------------------------
# load aggregate data
#-------------------------------------------------------------
dataDir  = "$(pwd())/data/Para" * nDataSel *"/"
agg_data = loadaggdata(1,Tend)
period_agg = size(agg_data)[1]
n_agg    = size(agg_data)[2]

#-------------------------------------------------------------
# Load coefficients from density estimation
#-------------------------------------------------------------
sNameLoadDir = "fVAR" * nfVARSpec
loaddir  = "$(pwd())/results/Para" *nDataSel * "/" * sNameLoadDir *"/";

PhatDensCoef_factor, MDD_GoF, VinvLam_all, PhatDensCoef_lambda, PhatDensCoef_mean  = loaddensdata(1,Tend,K,nfVARSpec)
n_cross = size(PhatDensCoef_factor)[2]

knots_all = readdlm(loaddir * "fVAR" *nfVARSpec * "_knots_all.csv", ',', Float64, '\n'; skipstart = 1);
ii=getindex.(findall(K_vec.-K.==0),[1 2])[1] # find index ii where K==K_vec
knots = knots_all[quant_sel[ii,:].==1]

# knots = loadknots(quant_vec, quant_sel)

#-------------------------------------------------------------
# Load posterior means
#-------------------------------------------------------------

sName    = "fVAR" * nfVARSpec * "_" * modName * nModSpec * "_" * "MCMC" * nMCMCSpec;
loadDir  = "$(pwd())/results/Para" *nDataSel * "/" * sName *"/";

Bpdraw     = load(loadDir * sName * "_PostDraws.jld", "Bpdraw")
#n_subseq                 = floor(Int,size(PHIpdraw)[1]/n_every)

YY_IRF = CSV.read(loadDir * sName * "_IRF_YY_" *nShockspec * "_pmean.csv", DataFrame, header = true);
PhatDens_IRF = CSV.read(loadDir * sName * "_IRF_PhatDens_" *nShockspec*"_pmean.csv", DataFrame, header = true);
YY_IRF = convert(Array, YY_IRF)
PhatDens_IRF = convert(Array, PhatDens_IRF)

#-------------------------------------------------------------
# Generate transformed IRFs at posterior mean
#-------------------------------------------------------------

println("")
println("Generating transformed IRFs at posterior mean... ")
time_init_loop = time_ns();

emp_percs_all    = zeros(H+1,length(vec_percs))
densvalues_ss    = PhatDens_IRF[1,:]; # steady state density
PhatMassDiff_IRF = zeros(H+1,1);
Gini_IRF         = zeros(H+1,1);

for hh = 1:(H+1)
    # read estimated states in period t, 1:n_agg are agg data, n_agg+1:end are dens coeff
    statepmean_hh = YY_IRF[hh,:]

    PhatDensCoef_hh    = coefRecover(statepmean_hh[n_agg+1:end]', PhatDensCoef_lambda, PhatDensCoef_mean)
    emp_percs_all[hh,:] = DensPercentiles(PhatDensCoef_hh, knots, 0.0, xgrid, grid_temp, vec_percs)
    PhatMassDiff_IRF[hh,1] = ProbMassDiff_Cutoff(PhatDens_IRF[hh,:], densvalues_ss, cutoff, xgrid);
    Gini_IRF[hh] = GiniCoef(PhatDens_IRF[hh,:], xgrid)

end

savedir = "$(pwd())/results/Para" *nDataSel * "/" * sName *"/";
CSV.write(savedir * sName * "_IRF_Pctl_" *nShockspec*"_pmean.csv", DataFrame(emp_percs_all))
CSV.write(savedir * sName * "_IRF_BelowCutoff_" *nShockspec*"_pmean.csv", DataFrame(PhatMassDiff_IRF))
CSV.write(savedir * sName * "_IRF_Gini_" *nShockspec*"_pmean.csv", DataFrame(Gini_IRF))

time_loop=signed(time_ns()-time_init_loop)/1000000000
println("Elapsed time = $(time_loop)")
println("")


#-------------------------------------------------------------
# Generate IRFs for a subset of posterior draws
#-------------------------------------------------------------

println("")
println("Generating transformed IRFs for posterior draws... ")
println("")

YY_IRF_uncertainty        = zeros(H+1,n_agg+n_cross,n_drawsread)
PhatDens_IRF_uncertainty  = zeros(H+1,length(xgrid),n_drawsread)

for pp = 1:n_drawsread
    YY_IRF_pp = CSV.read(loadDir * sName * "_IRF_YY_" *nShockspec * "_"*string(pp) *".csv", DataFrame, header = true);
    PhatDens_IRF_pp = CSV.read(loadDir * sName * "_IRF_PhatDens_" *nShockspec* "_" * string(pp) * ".csv", DataFrame, header = true);

    YY_IRF_uncertainty[:,:,pp] = convert(Array, YY_IRF_pp)
    PhatDens_IRF_uncertainty[:,:,pp] = convert(Array, PhatDens_IRF_pp)
end

emp_percs_uncertainty        = zeros(H+1,length(vec_percs),n_drawsread)
PhatMassDiff_IRF_uncertainty = zeros(H+1,n_drawsread);
Gini_IRF_uncertainty         = zeros(H+1,n_drawsread);

for pp = 1:n_drawsread

    time_init_loop = time_ns();
    println("Draw number:  $pp")
    println("Remaining draws:  $(n_drawsread-pp)")

    for hh = 1:(H+1)
        # read estimated states in period t, 1:n_agg are agg data, n_agg+1:end are dens coeff
        statepmean_hh = YY_IRF_uncertainty[hh,:,pp]

        PhatDensCoef_hh    = coefRecover(statepmean_hh[n_agg+1:end]', PhatDensCoef_lambda, PhatDensCoef_mean)
        emp_percs_uncertainty[hh,:,pp] = DensPercentiles(PhatDensCoef_hh, knots, 0.0, xgrid, grid_temp, vec_percs)
        PhatMassDiff_IRF_uncertainty[hh,pp] = ProbMassDiff_Cutoff(PhatDens_IRF_uncertainty[hh,:,pp], densvalues_ss, cutoff, xgrid);
        Gini_IRF_uncertainty[hh,pp] = GiniCoef(PhatDens_IRF_uncertainty[hh,:,pp], xgrid)
    end

    time_loop=signed(time_ns()-time_init_loop)/1000000000
    println("Elapsed time = $(time_loop)")
    println("")
    CSV.write(savedir * sName * "_IRF_Pctl_" *nShockspec * "_" * string(pp)* ".csv", DataFrame(emp_percs_uncertainty[:,:,pp]))

end

CSV.write(savedir * sName * "_IRF_BelowCutoff_" *nShockspec*"_uncertainty.csv", DataFrame(PhatMassDiff_IRF_uncertainty))
CSV.write(savedir * sName * "_IRF_Gini_" *nShockspec*"_uncertainty.csv", DataFrame(Gini_IRF_uncertainty))

##################################
# comparison to Winberry IRFs
##################################

#println("")
#println("Loading inputs... ")
#println("")

#dataDir = "$(pwd())/data/"
#simul_data       = CSV.read(dataDir * "simul_data_final_Tlong.csv", DataFrame, header = true); # run with version 1.5.
#densdraws_data   = CSV.read(dataDir * "simul_asset_data_final_Tlong.csv", DataFrame, header = false);
#employdraws_data = CSV.read(dataDir * "simul_employ_data_final_Tlong.csv", DataFrame, header = false);

#simul_data  = convert(Array,simul_data)
#densdraws   = convert(Array,densdraws_data)
#employdraws = convert(Array,employdraws_data)

# subsequently use the same knots regardless of sample size N,T
#knots_all = quantile(densdraws[employdraws.==1], quant_vec)

#Zt       = simul_data[:,1];
#agg_data = Zt[1:Tend,:]
#n_agg    = size(agg_data)[2]

#println("")
#println("Generating IRFs for KS-VAR... ")
#println("")

#measureCoef_2_1 = simul_data[1:Tend,8];
#measureCoef_2_2 = simul_data[1:Tend,9];
#measureCoef_2_3 = simul_data[1:Tend,10];

#moment_2_1 = simul_data[1:Tend,16];
#moment_2_2 = simul_data[1:Tend,17];
#moment_2_3 = simul_data[1:Tend,18];

#PDensCoef = [measureCoef_2_1 measureCoef_2_2 measureCoef_2_3 moment_2_1 moment_2_2 moment_2_3]
#(PDensCoef_factor, PDensCoef_lambda, PDensCoef_mean) = coefCompress(PDensCoef);

#aggEmployment = 0.93 # from Winberry's calibration


# load IRFs from script_ImpulseResponses.jl
#YY_IRF = CSV.read(loadDir * sName * "_PYY_IRF.csv", DataFrame, header = true);
#PDens_IRF = CSV.read(loadDir * sName * "_PDens_IRF.csv", DataFrame, header = true);
#YY_IRF = convert(Array, YY_IRF)
#PDens_IRF = convert(Array, PDens_IRF)

#println("")
#println("Generating transformed IRFs from Winberry... ")
#time_init_loop = time_ns();

#Pemp_percs_all    = zeros(H+1,length(vec_percs))
#densvalues_ss    = PDens_IRF[1,:]; # steady state density
#PMassDiff_IRF = zeros(H+1,1);
#PGini_IRF         = zeros(H+1,1);

#for hh = 1:(H+1)
    # read estimated states in period t, 1:n_agg are agg data, n_agg+1:end are dens coeff
#    statepmean_hh = YY_IRF[hh,:]
    # statepmean contains dev of unemployment from its mean -> add mean to obtain levels
#    unrate_hh = 0.0#statepmean_hh[3] + mean_unrate

#    PDensCoef_hh    = coefRecover(statepmean_hh[n_agg+1:end]', PDensCoef_lambda, PDensCoef_mean)
#    Pemp_percs_all[hh,:] = DensPercentilesWinberry(PDensCoef_hh, knots, aggEmployment, xgrid, grid_temp, vec_percs)
#    PMassDiff_IRF[hh,1] = ProbMassDiff_Cutoff(PDens_IRF[hh,:], densvalues_ss, cutoff, xgrid)
#    PGini_IRF[hh] = GiniCoef(PDens_IRF[hh,:], xgrid)

#end

#savedir = "$(pwd())/results/" * sName *"/";
#CSV.write(savedir * sName * "_IRF_Pctl_" *nShockspec*"_Winberry.csv", DataFrame(Pemp_percs_all))
#CSV.write(savedir * sName * "_IRF_BelowCutoff_" *nShockspec*"_Winberry.csv", DataFrame(PMassDiff_IRF))
#CSV.write(savedir * sName * "_IRF_Gini_" *nShockspec*"_Winberry.csv", DataFrame(PGini_IRF))
